#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "huffman.h"

// ================= 动态存储结构 =================
typedef struct {
    DTree node;
    int used;
} DNodeArr;

static DTree createDNode(char ch, int w) {
    DTree t = (DTree)malloc(sizeof(DNode));
    t->ch = ch;
    t->weight = w;
    t->left = t->right = NULL;
    return t;
}

DTree buildDTree(char chs[], int w[], int n) {
    DNodeArr arr[MAX_TREE_SIZE];
    for (int i = 0; i < n; ++i) {
        arr[i].node = createDNode(chs[i], w[i]);
        arr[i].used = 0;
    }
    int m = n;
    while (1) {
        int min1 = -1, min2 = -1;
        for (int i = 0; i < m; ++i) {
            if (!arr[i].used && (min1 == -1 || arr[i].node->weight < arr[min1].node->weight))
                min1 = i;
        }
        for (int i = 0; i < m; ++i) {
            if (!arr[i].used && i != min1 && (min2 == -1 || arr[i].node->weight < arr[min2].node->weight))
                min2 = i;
        }
        if (min2 == -1) break;
        DTree parent = createDNode('\0', arr[min1].node->weight + arr[min2].node->weight);
        parent->left = arr[min1].node;
        parent->right = arr[min2].node;
        arr[min1].used = arr[min2].used = 1;
        arr[m].node = parent;
        arr[m].used = 0;
        m++;
    }
    for (int i = 0; i < m; ++i)
        if (!arr[i].used) return arr[i].node;
    return NULL;
}

static void getDCodesRec(DTree root, char *code, int depth, CodeTable table[], int *idx) {
    if (!root) return;
    if (!root->left && !root->right) {
        table[*idx].ch = root->ch;
        code[depth] = '\0';
        strcpy(table[*idx].code, code);
        (*idx)++;
        return;
    }
    if (root->left) {
        code[depth] = '0';
        getDCodesRec(root->left, code, depth + 1, table, idx);
    }
    if (root->right) {
        code[depth] = '1';
        getDCodesRec(root->right, code, depth + 1, table, idx);
    }
}

void getDCodes(DTree root, CodeTable table[], int n) {
    char code[MAX_CODE_LEN];
    int idx = 0;
    getDCodesRec(root, code, 0, table, &idx);
}

// 优化版哈夫曼树可视化输出
void printDTreeInner(DTree root, char *prefix, int isLeft, int isLast) {
    if (!root) return;
    printf("%s", prefix);
    if (isLast) {
        printf("└──");
        strcat(prefix, "   ");
    } else {
        printf("├──");
        strcat(prefix, "│  ");
    }
    if (root->ch)
        printf("%c(%d)\n", root->ch == ' ' ? '_' : root->ch, root->weight);
    else
        printf("#(%d)\n", root->weight);

    int hasLeft = root->left != NULL;
    int hasRight = root->right != NULL;
    if (hasLeft || hasRight) {
        char nextPrefix[256];
        strcpy(nextPrefix, prefix);
        if (hasLeft)
            printDTreeInner(root->left, nextPrefix, 1, !hasRight);
        if (hasRight)
            printDTreeInner(root->right, nextPrefix, 0, 1);
    }
    // 恢复prefix
    prefix[strlen(prefix) - 3] = '\0';
}

void printDTree(DTree root, int depth) {
    char prefix[256] = "";
    printDTreeInner(root, prefix, 0, 1);
}

void freeDTree(DTree root) {
    if (!root) return;
    freeDTree(root->left);
    freeDTree(root->right);
    free(root);
}

static void findDCode(DTree root, char ch, char *code, int depth, char *res) {
    if (!root) return;
    if (!root->left && !root->right && root->ch == ch) {
        code[depth] = '\0';
        strcpy(res, code);
        return;
    }
    if (root->left) {
        code[depth] = '0';
        findDCode(root->left, ch, code, depth + 1, res);
    }
    if (root->right) {
        code[depth] = '1';
        findDCode(root->right, ch, code, depth + 1, res);
    }
}

void encodeD(DTree root, char *str, char *res) {
    char code[MAX_CODE_LEN];
    res[0] = '\0';
    for (int i = 0; str[i]; ++i) {
        char tmp[MAX_CODE_LEN] = "";
        findDCode(root, str[i], code, 0, tmp);
        strcat(res, tmp);
    }
}

void decodeD(DTree root, char *code, char *res) {
    int idx = 0;
    DTree p = root;
    for (int i = 0; code[i]; ++i) {
        if (code[i] == '0') p = p->left;
        else p = p->right;
        if (!p->left && !p->right) {
            res[idx++] = p->ch;
            p = root;
        }
    }
    res[idx] = '\0';
}

// ================= 静态存储结构 =================
int buildSTree(char chs[], int w[], int n, STree tree) {
    int m = 2 * n - 1;
    for (int i = 0; i < m; ++i) {
        tree[i].ch = (i < n) ? chs[i] : 0;
        tree[i].weight = (i < n) ? w[i] : 0;
        tree[i].parent = tree[i].left = tree[i].right = -1;
    }
    for (int i = n; i < m; ++i) {
        int min1 = -1, min2 = -1;
        for (int j = 0; j < i; ++j) {
            if (tree[j].parent == -1 && (min1 == -1 || tree[j].weight < tree[min1].weight))
                min1 = j;
        }
        for (int j = 0; j < i; ++j) {
            if (tree[j].parent == -1 && j != min1 && (min2 == -1 || tree[j].weight < tree[min2].weight))
                min2 = j;
        }
        tree[min1].parent = tree[min2].parent = i;
        tree[i].left = min1;
        tree[i].right = min2;
        tree[i].weight = tree[min1].weight + tree[min2].weight;
    }
    return m - 1; // 返回根节点下标
}

void getSCodes(STree tree, int n, CodeTable table[]) {
    char code[MAX_CODE_LEN];
    for (int i = 0; i < n; ++i) {
        int p = i, c = 0;
        int stack[MAX_CODE_LEN], top = 0;
        while (tree[p].parent != -1) {
            int parent = tree[p].parent;
            if (tree[parent].left == p) stack[top++] = 0;
            else stack[top++] = 1;
            p = parent;
        }
        table[i].ch = tree[i].ch;
        for (int j = 0; j < top; ++j)
            code[j] = stack[top - j - 1] + '0';
        code[top] = '\0';
        strcpy(table[i].code, code);
    }
}

void printSTree(STree tree, int n, int root, int depth) {
    if (root == -1) return;
    printSTree(tree, n, tree[root].right, depth + 1);
    for (int i = 0; i < depth; ++i) printf("    ");
    if (tree[root].ch) printf("%c(%d)\n", tree[root].ch, tree[root].weight);
    else printf("#(%d)\n", tree[root].weight);
    printSTree(tree, n, tree[root].left, depth + 1);
}

static void findSCode(STree tree, int n, char ch, char *res) {
    int idx = -1;
    for (int i = 0; i < n; ++i)
        if (tree[i].ch == ch) { idx = i; break; }
    if (idx == -1) { res[0] = '\0'; return; }
    int stack[MAX_CODE_LEN], top = 0, p = idx;
    while (tree[p].parent != -1) {
        int parent = tree[p].parent;
        if (tree[parent].left == p) stack[top++] = 0;
        else stack[top++] = 1;
        p = parent;
    }
    for (int j = 0; j < top; ++j)
        res[j] = stack[top - j - 1] + '0';
    res[top] = '\0';
}

void encodeS(STree tree, int n, char *str, char *res) {
    char code[MAX_CODE_LEN];
    res[0] = '\0';
    for (int i = 0; str[i]; ++i) {
        char tmp[MAX_CODE_LEN] = "";
        findSCode(tree, n, str[i], tmp);
        strcat(res, tmp);
    }
}

void decodeS(STree tree, int n, char *code, char *res) {
    int root = 2 * n - 2;
    int idx = 0, p = root;
    for (int i = 0; code[i]; ++i) {
        if (code[i] == '0') p = tree[p].left;
        else p = tree[p].right;
        if (tree[p].left == -1 && tree[p].right == -1) {
            res[idx++] = tree[p].ch;
            p = root;
        }
    }
    res[idx] = '\0';
}